import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.cm as cm
from skfem import *
from skfem.models.poisson import laplace, mass
from scipy.sparse.linalg import spsolve  # 添加稀疏矩阵求解器
# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

# 参数设置
l = 1.0          # 杆长
u0 = 1.0         # 初始温度常数
a2 = 0.1         # 增加热扩散系数以便更快看到变化
total_time = 2   # 减少总模拟时间
dt = 0.01        # 减小时间步长
n_elements = 100 # 增加单元数量以获得更平滑的结果

# 创建网格和函数空间
mesh = MeshLine(np.linspace(0, l, n_elements + 1))
basis = Basis(mesh, ElementLineP1())

# 标记边界
mesh = mesh.with_boundaries({
    'left': lambda x: x[0] == 0.0,
    'right': lambda x: x[0] == l
})

# 初始条件 - 使用更可靠的投影方法
def initial_temp(x):
    return u0 * x[0] / l

u_n = basis.project(initial_temp)  # 正确投影初始条件

# 准备矩阵: M + dt*a2*K
K = laplace.assemble(basis)  # 刚度矩阵 (∇u·∇v)
M = mass.assemble(basis)     # 质量矩阵 (u·v)
A = M + dt * a2 * K

# 边界条件处理
x0_dofs = basis.get_dofs(mesh.boundaries['left'])   # x=0 处边界
xl_dofs = basis.get_dofs(mesh.boundaries['right'])  # x=l 处边界

# 应用狄利克雷边界条件 (x=0, u=0) - 更可靠的方法
A, b_fixed = enforce(A, b=np.zeros(A.shape[0]), D=x0_dofs)

# 创建绘图和颜色映射
fig, ax = plt.subplots(figsize=(8, 4))
x_coords = mesh.p[0]

# 使用颜色表示温度
scatter = ax.scatter(x_coords, np.zeros_like(x_coords), c=u_n, cmap='hot', 
                    s=100, vmin=0, vmax=u0, marker='s', edgecolor='none')

# 添加颜色条
cbar = fig.colorbar(scatter, ax=ax)
cbar.set_label('温度', fontsize=12)
cbar.ax.tick_params(labelsize=10)

# 设置坐标轴
ax.set_xlim(0, l)
ax.set_ylim(-0.1, 0.1)  # 缩小y轴范围以突出颜色变化
ax.set_xlabel('位置 x', fontsize=12)
ax.set_title(f'细杆导热问题 (t = 0.0s)', fontsize=14)
ax.grid(True, alpha=0.3)

# 添加边界标记
ax.axvline(x=0, color='b', linestyle='--', linewidth=2, label='x=0 (固定为0℃)')
ax.axvline(x=l, color='g', linestyle='--', linewidth=2, label='x=l (绝热)')
ax.legend(loc='upper right', fontsize=10)

# 隐藏y轴刻度
ax.set_yticks([])

# 存储每一帧
frames = [u_n.copy()]

# 时间步进函数 - 修复求解过程
def update(frame):
    global u_n
    
    # 当前时间
    t = frame * dt
    
    # 构造右端向量
    b = M @ u_n
    
    # 应用狄利克雷边界条件 (x=0, u=0)
    b[x0_dofs] = 0.0  # 固定左端温度
    
    # 求解线性系统 - 使用稀疏矩阵求解器
    u_next = spsolve(A, b)
    
    # 确保右端绝热边界条件 (u_x=0)
    # 这通过自然边界条件自动满足，但我们可以显式处理
    # 实际上在有限元方法中，Neumann边界条件会自动处理
    
    # 更新温度场
    u_n = u_next
    
    # 更新图形
    scatter.set_array(u_n)
    ax.set_title(f'细杆导热问题 (t = {t:.2f}s)', fontsize=14)
    
    frames.append(u_n.copy())
    
    return scatter,

# 创建动画
ani = FuncAnimation(fig, update, frames=int(total_time/dt), 
                   interval=50, blit=True)

# 保存为GIF
ani.save('8-1例题2-细杆导热.gif', writer='pillow', fps=20)

plt.close()
print("计算完成！GIF 已保存为 '8-1例题2-细杆导热.gif'")